import os
import urllib.request
import zipfile
import sys
import shutil

def download_file(url, filename):
    print(f"正在下载: {filename}")
    try:
        urllib.request.urlretrieve(url, filename)
        print(f"下载成功: {filename}")
        return True
    except Exception as e:
        print(f"下载失败: {e}")
        return False

def unzip_file(zip_filename, extract_dir):
    print(f"正在解压: {zip_filename}")
    try:
        with zipfile.ZipFile(zip_filename, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)
        print(f"解压成功: {zip_filename}")
        return True
    except Exception as e:
        print(f"解压失败: {e}")
        return False

def find_model_files(start_dir):
    """递归查找模型文件"""
    model_files = []
    for root, dirs, files in os.walk(start_dir):
        for file in files:
            if file.endswith(".prototxt") or file.endswith(".caffemodel"):
                model_files.append(os.path.join(root, file))
    return model_files

def main():
    print("===== 性别和年龄模型下载脚本 v2 ======")

    # 模型文件列表
    models = [
        {
            "name": "gender_age_models",
            "url": "https://github.com/GilLevi/AgeGenderDeepLearning/archive/refs/heads/master.zip",
            "zip_file": "age_gender_models.zip",
            "extract_dir": "."
        }
    ]

    # 下载并解压模型
    for model in models:
        if not download_file(model["url"], model["zip_file"]):
            print(f"无法下载 {model['name']} 模型")
            continue

        if not unzip_file(model["zip_file"], model["extract_dir"]):
            print(f"无法解压 {model['name']} 模型")
            continue

        # 查找并移动模型文件
        try:
            # 查找解压后的目录
            extracted_dirs = [d for d in os.listdir(model["extract_dir"]) if os.path.isdir(d) and "AgeGenderDeepLearning" in d]
            if not extracted_dirs:
                print("未找到解压后的模型目录")
                continue

            src_dir = os.path.join(model["extract_dir"], extracted_dirs[0])
            print(f"找到模型目录: {src_dir}")

            # 递归查找模型文件
            model_files = find_model_files(src_dir)
            if not model_files:
                print("未找到模型文件")
                continue

            # 复制模型文件到当前目录
            for model_file in model_files:
                file_name = os.path.basename(model_file)
                dst_path = os.path.join(model["extract_dir"], file_name)
                shutil.copy2(model_file, dst_path)
                print(f"已复制: {file_name}")

            # 清理临时文件
            shutil.rmtree(src_dir)
            os.remove(model["zip_file"])
            print("已清理临时文件")
        except Exception as e:
            print(f"移动模型文件时出错: {e}")

    # 检查是否成功下载模型
    required_files = [
        "deploy_gender.prototxt",
        "gender_net.caffemodel",
        "deploy_age.prototxt",
        "age_net.caffemodel"
    ]

    missing_files = []
    for file in required_files:
        if not os.path.exists(file):
            missing_files.append(file)

    if not missing_files:
        print("\n===== 模型下载完成 ======")
        print("所有必需的模型文件已成功下载!")
    else:
        print("\n===== 模型下载不完全 ======")
        print("缺少以下文件:")
        for file in missing_files:
            print(f"- {file}")
        print("\n您可以尝试手动下载这些文件:")
        print("1. 访问: https://github.com/GilLevi/AgeGenderDeepLearning")
        print("2. 下载模型文件")
        print("3. 将文件放在当前目录")

if __name__ == "__main__":
    main()